from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from geoalchemy2 import Geometry
import psycopg2
from config import DATABASE_CONFIG
import logging

logger = logging.getLogger(__name__)

class DatabaseManager:
    """数据库管理类"""
    
    def __init__(self):
        self.config = DATABASE_CONFIG
        self.engine = None
        self.Session = None
        self._connect()
    
    def _connect(self):
        """建立数据库连接"""
        try:
            connection_string = (
                f"postgresql://{self.config['username']}:{self.config['password']}"
                f"@{self.config['host']}:{self.config['port']}/{self.config['database']}"
            )
            self.engine = create_engine(connection_string)
            self.Session = sessionmaker(bind=self.engine)
            logger.info("数据库连接成功")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise
    
    def get_session(self):
        """获取数据库会话"""
        return self.Session()
    
    def execute_query(self, query, params=None):
        """执行查询语句"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(query), params or {})
                return result.fetchall()
        except Exception as e:
            logger.error(f"查询执行失败: {e}")
            raise
    
    def test_connection(self):
        """测试数据库连接"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text("SELECT version()"))
                version = result.fetchone()[0]
                logger.info(f"数据库版本: {version}")
                
                # 检查PostGIS扩展
                result = conn.execute(text("SELECT PostGIS_Version()"))
                postgis_version = result.fetchone()[0]
                logger.info(f"PostGIS版本: {postgis_version}")
                return True
        except Exception as e:
            logger.error(f"数据库连接测试失败: {e}")
            return False

# 全局数据库管理器实例
db_manager = DatabaseManager() 